import SearchTree, { TreeNode } from "../07_二叉树/SearchTree";

export class AVLTreeNode<T> extends TreeNode<T> {
  left: AVLTreeNode<T> | null = null;
  right: AVLTreeNode<T> | null = null;
  root: AVLTreeNode<T> | null = null;

  private getHeight(): number {
    const leftHeight = this.left ? this.left.getHeight() : 0;
    const rightHeight = this.right ? this.right.getHeight() : 0;
    return Math.max(leftHeight, rightHeight) + 1;
  }

  private getBalanceFactor() {
    const leftHeight = this.left ? this.left.getHeight() : 0;
    const rightHeight = this.right ? this.right.getHeight() : 0;
    return leftHeight - rightHeight;
  }

  get isBalanced(): boolean {
    const factor = this.getBalanceFactor();
    // return factor >= -1 && factor <= 1;
    return Math.abs(factor) <= 1;
  }

  public get highestChild(): AVLTreeNode<T> | null {
    const leftHeight = this.left ? this.left.getHeight() : 0;
    const rightHeight = this.right ? this.right.getHeight() : 0;

    if (leftHeight > rightHeight) return this.left;
    if (leftHeight < rightHeight) return this.right;
    return this.isLeft ? this.left : this.right;
  }

  // 右旋
  rightRotation() {
    const isLeft = this.isLeft;
    const isRight = this.isRight;
    // 1. 处理pivot节点
    const pivot = this.left!;
    pivot.root = this.root;
    // 2. 处理pivot的right
    this.left = pivot.right;
    if (pivot.right) {
      pivot.right.root = this;
    }
    // 3. 处理this
    pivot.right = this;
    this.root = pivot;
    // 4. 挂载pivot
    if (!pivot.root) {
      // pivot直接作为tree的根
      return pivot;
    } else if (isLeft) {
      // pivot作为父节点的左子节点
      pivot.root.left = pivot;
    } else {
      // pivot作为父节点的右子节点
      pivot.root.right = pivot;
    }

    return pivot;
  }

  // 左旋
  lefeRotation() {
    const isLeft = this.isLeft;
    const isRight = this.isRight;
    const pivot = this.right!;
    pivot.root = this.root;

    this.right = pivot.left;
    if (pivot.left) {
      pivot.left.root = this;
    }

    pivot.left = this;
    this.root = pivot;

    if (!pivot.root) {
      return pivot;
    } else if (isLeft) {
      pivot.root.left = pivot;
    } else if (isRight) {
      pivot.root.right = pivot;
    }

    return pivot;
  }
}

export default class AVLTree<T> extends SearchTree<T> {
  // grant：找到不平衡的节点
  rebalance(root: AVLTreeNode<T>) {
    const pivot = root.highestChild;
    const current = pivot?.highestChild;

    let resultNode: AVLTreeNode<T> | null = null;
    if (pivot?.isLeft) {
      if (current?.isLeft) {
        // LL
        root.rightRotation();
      } else {
        // LR
        pivot.lefeRotation();
        resultNode = root.rightRotation();
      }
    } else {
      if (current?.isLeft) {
        // RL
        pivot?.rightRotation();
        resultNode = root.lefeRotation();
      } else {
        // RR
        root.lefeRotation();
      }
    }

    // 判断返回的pivot是否有父节点
    if (!resultNode?.root) {
      this.root = resultNode;
    }
  }

  protected createNode(value: T): TreeNode<T> {
    return new AVLTreeNode(value);
  }

  checkBalance(node: AVLTreeNode<T>) {
    let current = node.root;
    while (current) {
      if (!current.isBalanced) {
        this.rebalance(current);
      }
      current = current.root;
    }
  }
}

const avlTree = new AVLTree();
avlTree.insert(50);
avlTree.insert(100);
avlTree.insert(150);
avlTree.print();
